/*
 *   Copyright (C) 2014  Raghavan Santhanam, raghavanil4m@gmail.com.
 *   I wrote this cross-platform talk tool, WindowsLinuxTalk, for the
 *   love of programming entirely under Ubuntu(Linux) for Windows in
 *   a remote fashion. Send any opinions/suggestions/bug reports to
 *   raghavanil4m@gmail.com.
 *
 *   WindowsExec.c is free software: you can redistribute it and/or modify
 *   it under the terms of the GNU General Public License as published by
 *   the Free Software Foundation, either version 3 of the License, or
 *   (at your option) any later version.
 *
 *   WindowsExec.c is distributed in the hope that it will be useful,
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *   GNU General Public License for more details.
 *
 *   You should have received a copy of the GNU General Public License
 *   along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

/* WindowsExec.c -- Classic socket server that receives commands from Linux
 * client, intended to be executed under Windows and stream(send) back the
 * execution output(stdout and stderr) to Linux along with the exit code of
 * the command executed. */

#include <stdio.h>
#include <string.h>
#include <stdbool.h>
#include <setjmp.h>
#include <signal.h>
#include <errno.h>
#include <unistd.h>
#include <limits.h>
#include <stdlib.h>
#include <assert.h>

#include <winsock2.h>
#include <ws2tcpip.h>
#include <winbase.h>
#include <tlhelp32.h>

#include <fcntl.h>

#include <io.h>
_CRTIMP int __cdecl _chdir (const char*);
_CRTIMP char* __cdecl _getcwd (char*, int);
_CRTIMP int __cdecl _mkdir (const char*);

//#define __DEBUG__

#ifdef __DEBUG__
#	define LOG(...) do {\
		   			   	 printf(__VA_ARGS__);\
						 FAST_LOG(__VA_ARGS__);\
                    } while(0)
#else
#	define LOG(...)
#endif

//#define FAST_LOG_NEEDED
#ifdef FAST_LOG_NEEDED
static FILE *FastLogFp;
#	define START_FAST_LOG	do {\
                                 (void)remove("c:\\WindowsLinuxTalk\\FastLog-prev.txt");\
                                 (void)rename("c:\\WindowsLinuxTalk\\FastLog.txt", "c:\\WindowsLinuxTalk\\FastLog-prev.txt");\
		   					   	 (void)(!FastLogFp && (FastLogFp = fopen("c:\\WindowsLinuxTalk\\FastLog.txt", "a")));\
                            } while(0)
#	define FAST_LOG(...)	(void)(fprintf(FastLogFp, __VA_ARGS__) && fflush(FastLogFp))
#	define STOP_FAST_LOG	(void)(FastLogFp && fclose(FastLogFp) && (FastLogFp = NULL))
#else
#	define START_FAST_LOG
#	define FAST_LOG(...)
#	define STOP_FAST_LOG
#endif

static SOCKET WindowsServerSockfd = INVALID_SOCKET;
static struct sockaddr_in WindowsServAddrPort;
static const int WinPort = 7000;
static int WindowsRecvFd = -1;

static SOCKET WindowsServerKillSockfd = INVALID_SOCKET;
static struct sockaddr_in WindowsServAddrKillPort;
static const int WinKillPort = 7001;
static int WindowsRecvKillFd = -1;
static DWORD KillerThreadId = -1;

static FILE *StdoutFp;
static FILE *StderrFp;

static jmp_buf SaneState;

static void PrintProgName(char *n)
{
	const char *progName = n;
	char p[strlen(progName) + 1];
	memset(p, 0, sizeof(p));
	strcpy(p, progName);

	char *shortenName = p + strlen(p) - 1;
	while (shortenName >= p && *shortenName != '\\') {
		shortenName--;
	}
	
	if (shortenName < p || *shortenName == '\\') {
		shortenName++;
	}
	

	printf("\n\n ** %s **\n\n", shortenName);
}

static void KillLaunchedCommandCompletely(DWORD parentProcID);

static void CatchSignal(int sig)
{
	FAST_LOG("Caught : ");

    KillLaunchedCommandCompletely(GetCurrentProcessId());
    STOP_FAST_LOG;
    (void)remove("c:\\WindowsLinuxTalk\\FastLog-prev.txt");
    (void)rename("c:\\WindowsLinuxTalk\\FastLog.txt", "c:\\WindowsLinuxTalk\\FastLog-prev.txt");

	switch (sig) {
	case SIGSEGV:	LOG("SIGSEGV");
		 	if (StdoutFp) {
			   fclose(StdoutFp);
			   (void)rename("c:\\WindowsLinuxTalk\\stdout.txt", "c:\\WindowsLinuxTalk\\stdout_sigegv.txt");
			}   			   		
		 	if (StderrFp) {
			   fclose(StderrFp);
			   (void)rename("c:\\WindowsLinuxTalk\\stderr.txt", "c:\\WindowsLinuxTalk\\stderr_sigegv.txt");
			}   
			exit(0);			   					
			break;
	case SIGINT:	LOG("SIGINT\n");
			exit(0);
			break;
	case SIGTERM: 	LOG("SIGTERM");
            exit(0);       
			break;
	default:	LOG("Unknown/Some other signal(%d)", sig);
	        exit(0);  
			break;
	}

	LOG("\n");

	longjmp(SaneState, 1);
}

static void RegisterSignalHandlers(void)
{
	LOG("Registering signal handlers . . .\n");

	(void)signal(SIGSEGV, CatchSignal);
	(void)signal(SIGINT, CatchSignal);
	(void)signal(SIGTERM, CatchSignal);

	LOG("Registered.\n");
}

static void SaveState(void)
{
	LOG("Saving state . . .\n");

	bool firstTime = setjmp(SaneState) == 0;
	if (firstTime) {
		LOG("Saved.\n");
	} else {
		printf("Seems something wrong happened. Restoring to previous state . . .\n");
		RegisterSignalHandlers();
	}
}

static void Cleanup(void)
{
	if (WindowsRecvFd != -1) {
        close(WindowsRecvFd);
        WindowsRecvFd = -1;
    }

	if (WindowsServerSockfd != INVALID_SOCKET) {
		closesocket(WindowsServerSockfd);
		WindowsServerSockfd = INVALID_SOCKET;
	}
}

static bool PrepareForExecutingCmds(void)
{
	bool prepared = false;

	//LOG("Preparing for executing commands got from Linux . . .\n");

    Cleanup();

	WSADATA wsaData = { 0 };
	int iResult = WSAStartup(MAKEWORD(2, 2), &wsaData);
	if (iResult != NO_ERROR) {
		printf("Error: failed to WSAStartup - %d\n", iResult);
		goto finish;
	} 

	WindowsServerSockfd = socket(AF_INET, SOCK_STREAM, 0);
	bool opened = WindowsServerSockfd != INVALID_SOCKET;
	if (!opened) {
		printf("Error: failed to open socket - %ld\n", GetLastError());
		goto finish;
	}

	BOOL yes = 1;
	bool reuse = setsockopt(WindowsServerSockfd, SOL_SOCKET, SO_REUSEADDR, (char *)&yes, sizeof(yes)) != -1;
	if (!reuse) {
		printf("Error: failed to reuse socket - %s\n", strerror(errno));
		goto finish;
	}

	WindowsServAddrPort.sin_family = AF_INET;
	WindowsServAddrPort.sin_addr.s_addr = htonl(INADDR_ANY);
	WindowsServAddrPort.sin_port = htons(WinPort);

	bool bound = bind(WindowsServerSockfd, (struct sockaddr *)&WindowsServAddrPort, sizeof(struct sockaddr)) != -1;
	if (!bound) {
		printf("Error: failed to bind - %s\n", strerror(errno));
		goto finish;
	}

	bool listening = listen(WindowsServerSockfd, 10) != -1;
	if (!listening) {
		printf("Error: failed to listen - %s\n", strerror(errno));
		goto finish;
	}

	int s = sizeof(struct sockaddr);
	WindowsRecvFd = accept(WindowsServerSockfd, (struct sockaddr *)&WindowsServAddrPort, &s);
	bool accepted = WindowsRecvFd != -1;
	if (!accepted) {
		printf("Error: failed to accept - %s\n", strerror(errno));
		goto finish;
	}

	prepared = true;

	//LOG("Prepared.\n");

finish:

	if (!prepared) {
		Cleanup();
	}

	return prepared;
}

static void RefreshFp(FILE *fp)
{
 	long curOff = ftell(fp);
    if (curOff && feof(fp)) {			 
	   long newOff = fseek(fp, -1L, SEEK_CUR);
	   if (newOff < curOff) {
       	  (void)fseek(fp, 1L, SEEK_CUR);                              
	   }	  
    }
}

static bool SendData(FILE *standardStreamFp, FILE *dumpFileFp)
{
 	bool sent = true;   
 	   
//    fseek(standardStreamFp, 0L, SEEK_END);
//    fseek(dumpFileFp, 0L, SEEK_SET);
    long dumpedStdBytes = ftell(dumpFileFp);
    long stdBytes = ftell(standardStreamFp);
    bool anything = dumpedStdBytes < stdBytes;

    if (anything) {
        RefreshFp(dumpFileFp);          
        
        while (!feof(dumpFileFp) && !ferror(dumpFileFp)) {
            char s[1024] = "";
            size_t bytesRead = fread(s, 1, sizeof(s) - 1, dumpFileFp);
            bool fine = bytesRead > 0;
            
            if (fine) {
                int i = 0;
                while (i < bytesRead && s[i] == '\0') {
                      i++;
                }
                bool something = i != bytesRead;                                  
                
                if (something) { 
                    errno = 0;
                    sent = sendto(WindowsRecvFd, s + i, strlen(s + i), 0, NULL, 0) != -1;
                    if (!sent) {
                        FAST_LOG("Sending failed - %s\n", strerror(errno));
                        Cleanup();
                        break;
                    }
                }
            } else {
               sent = fine;		   		
               break;
            }
        }
    }
    
    return sent;
} 

static void EndSending(DWORD exitCode)
{
    // Say good bye!
    char s[1024] = "";
    sprintf(s, "$$ end(%ld) $$", exitCode);
    (void)sendto(WindowsRecvFd, s, strlen(s), 0, NULL, 0);
}

static void SendString(const char *s)
{
    (void)sendto(WindowsRecvFd, s, strlen(s), 0, NULL, 0);	   
}

static void RedirectStandardStreamsToFile(void)
{
 	remove("c:\\WindowsLinuxTalk\\stdout.txt");   
 	remove("c:\\WindowsLinuxTalk\\stderr.txt");
    (void)_mkdir("c:\\WindowsLinuxTalk\\");
    (void)_mkdir("c:\\WindowsLinuxTalk\\");     	
	StdoutFp = freopen("c:\\WindowsLinuxTalk\\stdout.txt", "w", stdout);	 	   
	StderrFp = freopen("c:\\WindowsLinuxTalk\\stderr.txt", "w", stderr);	
}

static void CloseRedirectedStreams(void)
{
 	if (StdoutFp) {
        fclose(StdoutFp);
        StdoutFp = NULL;
    }
    if (StderrFp) {
	    fclose(StderrFp);
	    StderrFp = NULL;
    }
}

static void EmptyFile(FILE *fp)
{
    (void)ftruncate(fileno(fp), 0);
    rewind(fp);
}

static void TrimLeadingSpaces(char **s)
{
 	char *t = *s;   
   	while (*t && *t == ' ') {
		  t++;
    }
    *s = t;
}     

static bool ExecuteCmds(void)
{
	bool executed = false;
	
	FAST_LOG("\n\nExecuting commands . . .\n");

	STARTUPINFO startupInfo = { 0 };
	PROCESS_INFORMATION procInfo = { 0 };

	FAST_LOG("cmd> ");
	char cmd[1024];
	memset(cmd, 0, sizeof(cmd));

	ssize_t bytesReceived = recvfrom(WindowsRecvFd, cmd, sizeof(cmd) - 1, 0, NULL, 0);
	bool received = bytesReceived > 0;
	if (!received) {
		if (bytesReceived == -1) {
		    FAST_LOG("Error: failed to recvfrom - %s\n", strerror(errno));
		} else {
			FAST_LOG("Seems like peer is dead. Resetting connection . . .\n");
            Cleanup();
		}
		goto bye;
	}

	FAST_LOG("%s\n", cmd);
	if (!strcmp(cmd, "exit")) {
		raise(SIGINT);
	}	

	if (!strncmp(cmd, "cd", 2) && cmd[2]) {
 	    char *dir = cmd + 2;
	   	TrimLeadingSpaces(&dir);
	   	
 	    bool changed = _chdir(dir) != -1;
 	    if (!changed) {
		   SendString(strerror(errno));
 		}
 	   
	    fflush(stderr);
        fflush(stdout);
		
		EndSending(0);
		
		goto done; 	   
    }
 
	char cmdArgs[sizeof("cmd.exe /C \"") - 1 + sizeof(cmd) - 1 + sizeof("\"")];
	memset(cmdArgs, 0, sizeof(cmdArgs));
	strcpy(cmdArgs, "cmd.exe /C \"");
	strcat(cmdArgs, cmd);
	strcat(cmdArgs, "\"");
	
	FAST_LOG("cmdArgs : %s\n", cmdArgs);

	EmptyFile(StdoutFp);
	EmptyFile(StderrFp);	
	
	FILE *dumpStdoutFp = fopen("c:\\WindowsLinuxTalk\\stdout.txt", "r");
	if (!dumpStdoutFp) {
 	   FAST_LOG("Error: failed to fopen - %s\n", strerror(errno));
 	   goto bye;
    }
    
 	FAST_LOG("%s() : %d\n", __func__, __LINE__);    

	FILE *dumpStderrFp = fopen("c:\\WindowsLinuxTalk\\stderr.txt", "r");
	if (!dumpStderrFp) {
 	   FAST_LOG("Error: failed to fopen - %s\n", strerror(errno));
 	   goto bye;
    }

	DWORD exitCode = -1;

	bool created = CreateProcess(NULL, cmdArgs, NULL, NULL, TRUE, NORMAL_PRIORITY_CLASS, NULL, NULL, &startupInfo, &procInfo);
	if (!created) {
		FAST_LOG("Error: failed to launch command(%s) - %ld\n", cmd, GetLastError());
	} else {
        while (1) {
 			  bool timedOut = WaitForSingleObject(procInfo.hProcess, 1) == WAIT_TIMEOUT;
 			  if (!timedOut) {     
                 RefreshFp(dumpStdoutFp);
                 RefreshFp(dumpStderrFp);                 
                             
		 		 (void)GetExitCodeProcess(procInfo.hProcess, &exitCode);

		 		 CloseHandle(procInfo.hThread);		
		 		 procInfo.hThread = INVALID_HANDLE_VALUE;
		 		 CloseHandle(procInfo.hProcess);			  	 			 
		 		 procInfo.hProcess = INVALID_HANDLE_VALUE;
 
	 			 break; // Ignore any error.
			  }
			  
			  fflush(StdoutFp);
			  fflush(StderrFp);
			  
	          bool sent = SendData(StdoutFp, dumpStdoutFp);
	          if (!sent) {
			  	 FAST_LOG("Sending failed -- 1\n");		 
			  	 
			  	 KillLaunchedCommandCompletely(GetCurrentProcessId());
			  	 exit(0);
			  	 // goto bye;
			  }	 
	          sent = SendData(StderrFp, dumpStderrFp); 
	          if (!sent) {
			  	 FAST_LOG("Sending failed -- 2\n");		 
				   			  	 		 		 		
                 KillLaunchedCommandCompletely(GetCurrentProcessId());	  	 
                 exit(0);
			  	 // goto bye;
			  }	 
			  
			  FAST_LOG("Sent.\n");
			  
			  Sleep(1);
        }        
        
		fflush(StdoutFp);
		fflush(StderrFp);
		
        bool sent = SendData(StdoutFp, dumpStdoutFp);         
        if (!sent) {
           FAST_LOG("Sending failed -- 3\n");		 			  	 		 
         		 		 			  	 
	       KillLaunchedCommandCompletely(GetCurrentProcessId());           
	       exit(0);
           // goto bye;
        }	 
 
      	sent = SendData(StderrFp, dumpStderrFp); 
        if (!sent) {
           FAST_LOG("Sending failed -- 4\n");		 			  	 		 
         	
           KillLaunchedCommandCompletely(GetCurrentProcessId());  	 		 			  	 
           exit(0);
           // goto bye;
        }
        
	    fflush(stderr);
        fflush(stdout);
		
		EndSending(exitCode);
		
		if (dumpStdoutFp) {
  		   fclose(dumpStdoutFp);
        }
        
		if (dumpStderrFp) {
  		   fclose(dumpStderrFp);
        }        
	}

	    fflush(stdout);
		fflush(stderr);
	
done:
	 
	 executed = true;
	 
	 FAST_LOG("\n\nExecuted.\n");
	 
bye:
	 return executed;
}

static void KillLaunchedCommandCompletely(DWORD parentProcID)
{      
       PROCESSENTRY32 pEntry = { .dwSize = sizeof(PROCESSENTRY32) };
       
       HANDLE snapHandle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
       bool snapped = snapHandle != INVALID_HANDLE_VALUE;
       if (!snapped) {
          FAST_LOG("Error: failed to CreateToolhelp32Snapshot - %ld\n", GetLastError());
          goto done;
       }

       bool someProc = Process32First(snapHandle, &pEntry) != ERROR_NO_MORE_FILES;
       if (someProc) {
          BOOL stillSomeProc = true;
          
          while (stillSomeProc) {
                bool child = pEntry.th32ParentProcessID == parentProcID;
                
                FAST_LOG("childParent : %ld parent : %ld\n", pEntry.th32ParentProcessID, parentProcID);
                
                if (child) {
                   DWORD childProcID = pEntry.th32ProcessID;
       
                   KillLaunchedCommandCompletely(childProcID);

                   HANDLE childHandle = OpenProcess(PROCESS_ALL_ACCESS, TRUE, childProcID);
                   if (childHandle) {
                      TerminateProcess(childHandle, -1);
                      CloseHandle(childHandle);
                      childHandle = NULL;
                   } else {
                      FAST_LOG("Error: failed to OpenProcess - %ld\n", GetLastError());
                   }
                }
                
                stillSomeProc = Process32Next(snapHandle, &pEntry);
                FAST_LOG("Next process : %s\n", pEntry.szExeFile);
          }
       }

done: ;
}

static bool PrepareForKillingCmds(void);

static DWORD WINAPI KillerThread(LPVOID args)
{
      while (1) {
            char s[sizeof("stop")] = "";
            ssize_t bytesReceived = recvfrom(WindowsRecvKillFd, s, sizeof(s) - 1, 0, NULL, 0);
            bool something = bytesReceived > 0;
            bool err = bytesReceived == SOCKET_ERROR;
            
            FAST_LOG("\n\n**bytesReceived : %ld**\n\n", bytesReceived);
            
            if (something) {
               bool stop = !strcmp(s, "stop");               
               if (stop) {
                  DWORD daddyProcID = GetCurrentProcessId();       
                  KillLaunchedCommandCompletely(daddyProcID);
               }
            } else {
              if (err) {     
                 FAST_LOG("Error: failed to recvfrom - %d\n", WSAGetLastError());
              }
     
              bool killPrepared = false;
              while (!killPrepared) {
                    killPrepared = PrepareForKillingCmds();
                    if (!killPrepared) {
                       FAST_LOG("Error: failed to kill-prepare.\n");
                    }
              }   
            }
      }
      
      return 0;
}

static void CleanupKill(void)
{
	if (WindowsRecvKillFd != -1) {
        close(WindowsRecvKillFd);
        WindowsRecvKillFd = -1;
    }

	if (WindowsServerKillSockfd != INVALID_SOCKET) {
		closesocket(WindowsServerKillSockfd);
		WindowsServerKillSockfd = INVALID_SOCKET;
	}
}

static bool PrepareForKillingCmds(void)
{
	bool prepared = false;

	//LOG("Preparing for killing commands got from Linux . . .\n");

    CleanupKill();

	WSADATA wsaData = { 0 };
	int iResult = WSAStartup(MAKEWORD(2, 2), &wsaData);
	if (iResult != NO_ERROR) {
		printf("Error: failed to WSAStartup - %d\n", iResult);
		goto finish;
	} 

	WindowsServerKillSockfd = socket(AF_INET, SOCK_STREAM, 0);
	bool opened = WindowsServerKillSockfd != INVALID_SOCKET;
	if (!opened) {
		printf("Error: failed to open socket - %ld\n", GetLastError());
		goto finish;
	}

	BOOL yes = 1;
	bool reuse = setsockopt(WindowsServerKillSockfd, SOL_SOCKET, SO_REUSEADDR, (char *)&yes, sizeof(yes)) != -1;
	if (!reuse) {
		printf("Error: failed to reuse socket - %s\n", strerror(errno));
		goto finish;
	}

	WindowsServAddrKillPort.sin_family = AF_INET;
	WindowsServAddrKillPort.sin_addr.s_addr = htonl(INADDR_ANY);
	WindowsServAddrKillPort.sin_port = htons(WinKillPort);

	bool bound = bind(WindowsServerKillSockfd, (struct sockaddr *)&WindowsServAddrKillPort, sizeof(struct sockaddr)) != -1;
	if (!bound) {
		printf("Error: failed to bind - %s\n", strerror(errno));
		goto finish;
	}

	bool listening = listen(WindowsServerKillSockfd, 10) != -1;
	if (!listening) {
		printf("Error: failed to listen - %s\n", strerror(errno));
		goto finish;
	}

	int s = sizeof(struct sockaddr);
	WindowsRecvKillFd = accept(WindowsServerKillSockfd, (struct sockaddr *)&WindowsServAddrKillPort, &s);
	bool accepted = WindowsRecvKillFd != -1;
	if (!accepted) {
		printf("Error: failed to accept - %s\n", strerror(errno));
		goto finish;
	}
	
	bool killerCreated = CreateThread(NULL, 0, KillerThread, NULL, 0, &KillerThreadId);
    if (!killerCreated) {
       printf("Error: failed to CreateThread - %ld\n", GetLastError());
       goto finish;
    }

	prepared = true;

	//LOG("Prepared.\n");

finish:

	if (!prepared) {
		CleanupKill();
	}

	return prepared;
}

int main(int argc, char *argv[])
{
 	START_FAST_LOG; 
 	
	PrintProgName(argv[0]);

	SaveState();

	RegisterSignalHandlers();

    bool killPrepared = false;
    while (!killPrepared) {
          killPrepared = PrepareForKillingCmds();
          if (!killPrepared) {
             FAST_LOG("Error: failed to kill-prepare.\n");
          }
    }

    bool redirected = false;

	while (1) {
		bool prepared = PrepareForExecutingCmds();
		if (!prepared) {
			FAST_LOG("Error: failed to prepare.\n");
			continue;
		}

        if (!redirected) {
		   printf("Ready.\n");				  
		   RedirectStandardStreamsToFile();
		   redirected = true;
		}

		while (1) {
			bool executed = ExecuteCmds();
			if (!executed) {
				FAST_LOG("Error: failed to execute.\n");
				break;
			}

			Sleep(10);
		}

		Sleep(10);
	}
	
	CloseRedirectedStreams();
		
	STOP_FAST_LOG;

	return 0;
}



